# Set the ORT_DYLIB to point to the onnxruntime shared lib.
# This should be somewhere in the target dir after running `maturin develop --release`

%env ORT_DYLIB_PATH=/home/adr/build/cargo/release/libonnxruntime.so

import pymc as pm
import pytensor
import pytensor.tensor as pt
import numpy as np
import torch
import io
import nutpie

N = 1000

# Create a model that represents the model and compute the logp and gradient
class LogpModel(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        #self.matrix = torch.nn.Parameter(torch.zeros(1_000, N), requires_grad=False)

    def forward(self, x, bias=None):
        #return -(x * x).sum() * 0.5 + (self.matrix @ x).sum(), -x
        return -(x * x).sum() * 0.5, -x

model = LogpModel()

# Export the model to onnx
x0 = torch.ones(N)
onnx_program = torch.onnx.dynamo_export(model, x0)
onnx_code = io.BytesIO()
onnx_program.save(onnx_code)
onnx_code = onnx_code.getvalue()

# Set up onnx options
providers = nutpie._lib.OnnxProviders()
providers.add_cuda()
providers.add_cpu()

# Create nutpie model and sample
model = nutpie._lib.OnnxModel(N, onnx_code, providers)
settings = nutpie._lib.PyNutsSettings.Diag()
progress = nutpie._lib.ProgressType.none()
sampler = nutpie._lib.PySampler.from_onnx(settings, 6, model, progress)
tr = sampler.wait()

%%time
progress = nutpie._lib.ProgressType.none()
sampler = nutpie._lib.PySampler.from_onnx(settings, 6, model, progress)
tr = sampler.wait()